import numpy as np

import k3d
from k3d.colormaps import paraview_color_maps
from k3d.headless import k3d_remote, get_headless_driver


def generate():
    plot = k3d.plot(grid_visible=False,
                    camera_auto_fit=False,
                    screenshot_scale=1.0,
                    axes_helper=0)

    bodies_count = 40

    bodies = np.array([[0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,
                        0.0000000e+00, 0.0000000e+00, 1.0000000e+06],
                       [-4.9679345e-01, 2.6570296e-01, 3.4227931e-01, -2.2651915e-02,
                        -2.5265724e-03, -9.3185175e-03, 6.3181415e+02],
                       [3.7021518e-02, -3.3532095e-01, -1.3006118e-01, 3.8169832e-03,
                        -1.7639263e-02, 1.7560443e-02, 1.2228839e+03],
                       [3.2473081e-01, -3.8958952e-01, 3.6065096e-01, -1.6591853e-02,
                        2.2979354e-02, 1.5027565e-02, 1.1379092e+03],
                       [1.9925886e-01, -4.3243438e-02, 2.3790616e-01, -2.4283929e-02,
                        2.2036737e-02, 1.5822455e-02, 7.3010986e+02],
                       [-4.4143024e-01, -7.1139961e-02, 1.7319494e-01, -1.1633939e-02,
                        5.0242455e-03, -5.0497951e-04, 6.7444598e+02],
                       [7.8688264e-02, -1.7087817e-01, -1.5081182e-01, -2.4320260e-03,
                        7.0646438e-03, -1.7999319e-02, 9.7800610e+02],
                       [-5.9594750e-02, -7.2615296e-02, 4.3256253e-01, -1.0310315e-02,
                        -1.0246425e-02, 1.0331392e-02, 5.5932916e+02],
                       [-2.2181371e-01, 1.2450159e-02, -3.5403877e-02, 2.4978900e-02,
                        -9.8985955e-03, -1.0434589e-02, 1.2175408e+03],
                       [2.1475327e-01, 4.1201133e-01, -3.1412601e-02, -1.1963538e-02,
                        1.0723472e-03, -3.5631433e-03, 1.1424460e+03],
                       [4.2195350e-01, 3.7973666e-01, -1.8134981e-02, -1.9921964e-02,
                        -7.1983635e-03, 7.6964349e-03, 1.4089318e+03],
                       [2.1617496e-01, -1.9794500e-01, -8.7558389e-02, -1.6835531e-02,
                        -1.5182316e-02, -1.3301859e-02, 5.5457037e+02],
                       [-4.1794181e-02, -1.9101912e-01, 1.8947464e-01, 1.4356637e-02,
                        -1.5364028e-02, 1.2453362e-02, 1.0287977e+03],
                       [2.2971267e-01, -2.8791130e-03, -2.1326435e-01, -7.8020454e-03,
                        2.5177689e-03, -9.7344117e-03, 1.0474579e+03],
                       [2.6785791e-02, 1.6838884e-01, 2.1399486e-01, 2.4175057e-02,
                        2.4402628e-02, -1.3445861e-02, 1.0599727e+03],
                       [3.1909162e-01, 3.1714398e-01, 2.0338118e-02, -2.4140883e-02,
                        1.5961260e-02, -2.0101408e-02, 9.4045221e+02],
                       [-3.3931443e-01, -2.2721624e-01, 9.7239554e-02, 3.6300363e-03,
                        -1.9278442e-02, 1.9151056e-02, 6.7030249e+02],
                       [6.6103995e-02, -4.8539740e-01, 1.7521322e-01, 1.8680961e-03,
                        9.0755373e-03, 1.3850155e-02, 9.4479297e+02],
                       [-1.6933033e-01, 5.7348192e-02, 2.0398575e-01, -1.0792896e-02,
                        -1.6141150e-02, 8.3085243e-03, 1.4722698e+03],
                       [3.5467851e-01, 2.5269783e-01, -4.7175428e-01, -2.1373475e-02,
                        4.0807934e-03, 2.4820814e-02, 8.4746240e+02],
                       [-2.7038407e-01, -4.8006400e-01, -3.2617545e-01, -2.1090398e-02,
                        1.7032122e-02, 1.7789362e-02, 6.2347046e+02],
                       [-4.8136929e-01, -3.3214843e-01, -1.8433192e-01, -8.0333194e-03,
                        -1.6524477e-02, 6.1281892e-03, 9.0508881e+02],
                       [5.4457009e-02, 1.7607331e-02, 6.7196190e-02, -1.7219210e-02,
                        1.9105632e-02, 2.4653585e-03, 5.7763861e+02],
                       [-1.3687915e-01, 8.4126353e-02, -1.3019252e-01, -1.8430030e-02,
                        2.0881360e-02, -2.4473562e-03, 5.5547150e+02],
                       [1.0026491e-01, 2.7884889e-01, -3.1183958e-03, 6.0113161e-03,
                        -1.4223135e-02, -1.8654644e-02, 5.5057245e+02],
                       [2.2953439e-01, -2.0373765e-01, 1.9156224e-01, -3.1318248e-03,
                        -2.5028258e-03, 8.5311802e-03, 6.9054175e+02],
                       [-4.9273950e-01, 2.5772792e-01, 4.5453101e-01, -1.1930334e-02,
                        -1.0238136e-02, 3.0950904e-03, 5.3005304e+02],
                       [3.5183370e-02, 4.1218692e-01, -3.8023493e-01, 1.6193924e-02,
                        8.6651091e-03, -5.8665574e-03, 8.8095959e+02],
                       [-3.0682594e-01, -1.1159393e-01, 3.8180280e-01, -1.8235832e-02,
                        2.0600855e-02, -5.5592731e-03, 1.2845226e+03],
                       [1.3181090e-01, -1.3926461e-01, 3.5556364e-01, 1.8209113e-02,
                        -7.5311526e-03, -2.0307558e-02, 9.4135608e+02],
                       [2.9862744e-01, 2.6593310e-01, 8.9805841e-02, 2.3497565e-02,
                        -2.3358732e-03, 6.1519952e-03, 6.3521173e+02],
                       [1.8002808e-02, -3.1908777e-01, -1.4509407e-01, 1.7548909e-02,
                        1.0903591e-02, -2.2118099e-02, 1.2389504e+03],
                       [-4.0650040e-02, -2.2764242e-01, -7.5831264e-02, -5.7892441e-03,
                        3.0585080e-03, -2.3829598e-02, 1.1236339e+03],
                       [-2.4570310e-01, 2.0286041e-01, 1.2090033e-01, -1.4608744e-02,
                        1.2567866e-02, -2.4784094e-02, 9.3537274e+02],
                       [3.9212245e-01, -1.8226206e-02, -3.4510234e-01, 2.7927459e-04,
                        -1.9444963e-02, 6.7013353e-03, 6.1505707e+02],
                       [-4.4596395e-01, -4.6554554e-01, -3.3184126e-01, 1.9790804e-02,
                        -2.9474200e-04, 1.0445378e-02, 6.7410992e+02],
                       [-4.7785598e-01, -4.4018179e-01, 3.1628680e-01, 3.2713532e-03,
                        1.2704453e-02, 1.4309299e-02, 7.2144165e+02],
                       [1.7605799e-01, 4.6154976e-01, 3.6913526e-01, -2.0895386e-02,
                        -2.3292074e-02, 4.9127610e-03, 6.6594043e+02],
                       [1.8955290e-02, -3.2265896e-01, 4.2756569e-01, -1.1221913e-02,
                        -2.0262299e-02, 1.1423305e-02, 9.2574188e+02],
                       [1.0150486e-01, -2.8043854e-01, -3.0282909e-01, 1.2297285e-02,
                        1.9709280e-02, 1.1570200e-02, 1.3365149e+03]]).astype(np.float32)

    for i in range(1, bodies_count):
        bodies[i, 0:3] = (
                                 bodies[i, 0:3] / np.linalg.norm(bodies[i, 0:3])) * 0.5

    points = k3d.points(bodies[:, 0:3],
                        point_size=0.03,
                        color=0x3e3a3a)
    plot += points

    G = 6.67E-11
    lines = []
    speeds = []
    positions = {}

    for i in range(bodies_count):
        lines.append([])
        speeds.append([])

    for t in range(500):
        for i in range(bodies_count):
            sum_force = np.zeros(3)

            for j in range(bodies_count):
                if i == j:
                    continue

                direction = bodies[j, 0:3] - bodies[i, 0:3]
                force = G * bodies[i, 6] * bodies[j, 6] * direction
                force = force / (np.linalg.norm(direction) ** 3)
                sum_force = sum_force + force

            bodies[i, 3:6] = bodies[i, 3:6] + sum_force / bodies[i, 6]

        for i in range(bodies_count):
            bodies[i, 0:3] = bodies[i, 0:3] + bodies[i, 3:6] * 0.15
            lines[i].append(np.copy(bodies[i, 0:3]))
            speeds[i].append(np.linalg.norm(bodies[i, 3:6]))

        positions[str(t * 0.01)] = np.copy(bodies[:, 0:3]).astype(np.float32)

    for line, speed in zip(lines, speeds):
        plot += k3d.line(np.array(line).astype(np.float32),
                         width=0.0002,
                         attribute=speed,
                         color_range=[0, 0.1],
                         color_map=paraview_color_maps.Erdc_iceFire_H)

    points.positions = positions

    plot.camera = [1.5491, -1.2661, -0.3120,
                   -0.1189, 0.0576, -0.1350,
                   0.6329, 0.7390, -0.2306]

    headless = k3d_remote(plot, get_headless_driver(), width=800, height=800)

    headless.sync(hold_until_refreshed=True)
    # headless.camera_reset(0.85)

    screenshot = headless.get_screenshot()
    headless.close()

    return screenshot
